'''
重返未来1999人物识别代码
'''

import cv2
import torch,os
import traceback
"""
torch_init推理初始化函数
传参说明
    device---推理设备(默认CPU)
    yolo_path---yolov5项目文件路径(在程序里需自行修改)
    model_path---.pt模型文件路径(在程序里需自行修改)
"""
def torch_init(device, yolo_path, model_path):
    device = torch.device(device)  # 设置推理设备

    torch_model = torch.hub.load(yolo_path, 'custom',
                                 model_path, source='local',
                                 force_reload=True)  # 加载本地yolov5模型(需要修改路径和文件)
    return device, torch_model  # 返回
"""
Get_Object获取物体对应下标函数
传参说明:
    number---识别物体下表
返回
    line--classes.txt里面查询到的第几行，即人物下标
"""
def Get_Object(number):
    classes = open('classes.txt',encoding = 'utf-8')#读取classes.txt文件
    
    for num, line in enumerate(classes):
        if num == number:
            break
    return line
"""
Detect_Enemy识别敌方函数
传参说明
    frame---传入图像
    size---输入图像大小(默认640*640)
    conference---置信度
    device---推理设备
    torch_model---模型
返回
    返回敌方的类型,坐标信息列表enemy_information

<--注:下面的函数基本上都是大同小异，会简易概括__>
"""
def Detect_Enemy(frame, size, conference, device, torch_model):
    enemy_information = []
    torch_model = torch_model.to(device)
    results = torch_model(frame, size=size)  # 推理图像
    
    try:  # 尝试
        xyxy = results.pandas().xyxy[0].values
        xmins, ymins, xmaxs, ymaxs, class_list, confidences = xyxy[:, 0], xyxy[:, 1],\
                                                                  xyxy[:, 2], xyxy[:, 3], xyxy[:,5], xyxy[:, 4]
        for xmin, ymin, xmax, ymax, class_l, conf in zip(xmins, ymins, xmaxs, ymaxs, class_list, confidences):
            if conf >= conference and (int(class_l) == 1 or int(class_l) == 2 or int(class_l) == 3):  # 如果置信度大于0.3
                enemy_information.append([int(xmin), int(ymin), int(xmax), int(ymax),int(class_l)])  # 将识别物体信息加入armor_list列表
    except:
        traceback.print_exc()  # 打印报错信息
        return enemy_information

    return enemy_information
"""
Detect_Character识别我方角色函数
"""
def Detect_Character(frame, size, conference, device, torch_model):
    character_information = []
    torch_model = torch_model.to(device)
    results = torch_model(frame, size=size)  # 推理图像
    
    try:  # 尝试
        xyxy = results.pandas().xyxy[0].values
        xmins, ymins, xmaxs, ymaxs, class_list, confidences = xyxy[:, 0], xyxy[:, 1],\
                                                                  xyxy[:, 2], xyxy[:, 3], xyxy[:,5], xyxy[:, 4]
        for xmin, ymin, xmax, ymax, class_l, conf in zip(xmins, ymins, xmaxs, ymaxs, class_list, confidences):
            if conf >= conference and (int(class_l) == 4 or int(class_l) == 5 or int(class_l) == 6):  # 如果置信度大于0.3
                character_information.append([int(xmin), int(ymin), int(xmax), int(ymax),int(class_l)])  # 将识别物体信息加入armor_list列表
    except:
        traceback.print_exc()  # 打印报错信息
        return character_information

    return character_information
"""
Detect_Cost识别司辰技能"最初的速度"函数
"""
def Detect_Cost(frame, size, device, conference,torch_model):
    torch_model = torch_model.to(device)
    results = torch_model(frame, size=size)  # 推理图像
    
    try:  # 尝试
        xyxy = results.pandas().xyxy[0].values
        xmins, ymins, xmaxs, ymaxs, class_list, confidences = xyxy[:, 0], xyxy[:, 1],\
                                                              xyxy[:, 2], xyxy[:, 3], xyxy[:,5], xyxy[:, 4]
        for xmin, ymin, xmax, ymax, class_l, conf in zip(xmins, ymins, xmaxs, ymaxs, class_list, confidences):
            if int(class_l) == 0 and conf >= conference:
                cv2.rectangle(frame,(int(xmin),int(ymin)),(int(xmax),int(ymax)),(255, 0, 225), 2)
                return True
                break
        return False

    except:  # 若出现错误：
        traceback.print_exc()  # 打印报错信息
        return False




